// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <nebula/core/extension_type.h>

namespace nebula {
    namespace extension {
        class TURBO_EXPORT FixedShapeTensorArray : public ExtensionArray {
        public:
            using ExtensionArray::ExtensionArray;

            /// \brief create a FixedShapeTensorArray from a Tensor
            ///
            /// This method will create a FixedShapeTensorArray from a Tensor, taking its first
            /// dimension as the number of elements in the resulting array and the remaining
            /// dimensions as the shape of the individual tensors. If Tensor provides strides,
            /// they will be used to determine dimension permutation. Otherwise, row-major layout
            /// (i.e. no permutation) will be assumed.
            ///
            /// \param[in] tensor The Tensor to convert to a FixedShapeTensorArray
            static turbo::Result<std::shared_ptr<FixedShapeTensorArray> > FromTensor(
                const std::shared_ptr<Tensor> &tensor);

            /// \brief create a Tensor from FixedShapeTensorArray
            ///
            /// This method will create a Tensor from a FixedShapeTensorArray, setting its first
            /// dimension as length equal to the FixedShapeTensorArray's length and the remaining
            /// dimensions as the FixedShapeTensorType's shape. Shape and dim_names will be
            /// permuted according to permutation stored in the FixedShapeTensorType metadata.
            const turbo::Result<std::shared_ptr<Tensor> > ToTensor() const;
        };

        /// \brief Concrete type class for constant-size Tensor data.
        /// This is a canonical nebula extension type.
        /// See: https://nebula.apache.org/docs/format/CanonicalExtensions.html
        class TURBO_EXPORT FixedShapeTensorType : public ExtensionType {
        public:
            FixedShapeTensorType(const std::shared_ptr<DataType> &value_type, const int32_t &size,
                                 const std::vector<int64_t> &shape,
                                 const std::vector<int64_t> &permutation = {},
                                 const std::vector<std::string> &dim_names = {})
                : ExtensionType(fixed_size_list(value_type, size)),
                  value_type_(value_type),
                  shape_(shape),
                  permutation_(permutation),
                  dim_names_(dim_names) {
            }

            std::string extension_name() const override { return "nebula.fixed_shape_tensor"; }

            std::string to_string(bool show_metadata = false) const override;

            /// Number of dimensions of tensor elements
            size_t ndim() const { return shape_.size(); }

            /// Shape of tensor elements
            const std::vector<int64_t> &shape() const { return shape_; }

            /// Value type of tensor elements
            const std::shared_ptr<DataType> &get_value_type() const { return value_type_; }

            /// Strides of tensor elements. Strides state offset in bytes between adjacent
            /// elements along each dimension. In case permutation is non-empty strides are
            /// computed from permuted tensor element's shape.
            const std::vector<int64_t> &strides();

            /// Permutation mapping from logical to physical memory layout of tensor elements
            const std::vector<int64_t> &permutation() const { return permutation_; }

            /// Dimension names of tensor elements. Dimensions are ordered physically.
            const std::vector<std::string> &dim_names() const { return dim_names_; }

            bool ExtensionEquals(const ExtensionType &other) const override;

            std::string serialize() const override;

            turbo::Result<std::shared_ptr<DataType> > deserialize(
                std::shared_ptr<DataType> storage_type,
                const std::string &serialized_data) const override;

            /// create a FixedShapeTensorArray from ArrayData
            std::shared_ptr<Array> make_array(std::shared_ptr<ArrayData> data) const override;

            /// \brief create a Tensor from an ExtensionScalar from a FixedShapeTensorArray
            ///
            /// This method will return a Tensor from ExtensionScalar with strides
            /// derived from shape and permutation of FixedShapeTensorType. Shape and
            /// dim_names will be permuted according to permutation stored in the
            /// FixedShapeTensorType metadata.
            static turbo::Result<std::shared_ptr<Tensor> > MakeTensor(
                const std::shared_ptr<ExtensionScalar> &scalar);

            /// \brief create a FixedShapeTensorType instance
            static turbo::Result<std::shared_ptr<DataType> > create(
                const std::shared_ptr<DataType> &value_type, const std::vector<int64_t> &shape,
                const std::vector<int64_t> &permutation = {},
                const std::vector<std::string> &dim_names = {});

        private:
            std::shared_ptr<DataType> storage_type_;
            std::shared_ptr<DataType> value_type_;
            std::vector<int64_t> shape_;
            std::vector<int64_t> strides_;
            std::vector<int64_t> permutation_;
            std::vector<std::string> dim_names_;
        };

        /// \brief Return a FixedShapeTensorType instance.
        TURBO_EXPORT std::shared_ptr<DataType> fixed_shape_tensor(
            const std::shared_ptr<DataType> &storage_type, const std::vector<int64_t> &shape,
            const std::vector<int64_t> &permutation = {},
            const std::vector<std::string> &dim_names = {});
    } // namespace extension
} // namespace nebula
